import torch
import torch.nn as nn


class FNN(nn.Module):
    def __init__(self, in_dim=100, out_dim=10):
        super(FNN, self).__init__()
        self.layer1 = nn.Linear(in_dim, 64)
        self.rule = nn.ReLU(inplace=True)
        self.layer2 = nn.Linear(64, out_dim)
        self.softmax = nn.Softmax(dim=0)
        for params in self.parameters():
            nn.init.normal_(params, mean=0, std=0.01)

    def forward(self, x):
        x = self.layer1(x)
        x = self.rule(x)
        x = self.layer2(x)
        x = self.softmax(x)
        return x


def train(train_loader, model, criterion, optimizer):
    model.train()
    for i, (feature, label) in enumerate(train_loader):
        output = model(feature)
        loss = criterion(output, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


def validate(val_loader, model, criterion):
    model.eval()
    correct = 0
    with torch.no_grad():
        for i, (feature, label) in enumerate(val_loader):
            output = model(feature)
            _, predict = output.topk(1, 1, True, True)
            correct += predict.eq(label.view(-1, 1)).sum(0, keepdim=True)
    return correct


if __name__ == '__main__':
    model = FNN(in_dim=100, out_dim=10)
    epoch = 20
    batch_size = 10
    lr = 0.2
    critierion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    train_points = torch.randn(50, 100)
    train_labels = torch.randint(0, 10, size=(50,))
    train_data = torch.utils.data.TensorDataset(train_points, train_labels)
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False)
    for i in range(epoch):
        train(train_loader=train_loader, model=model, criterion=critierion, optimizer=optimizer)
        correct = validate(val_loader, model, critierion)
        print("epoch:\t", i, "\tacc:\t", correct * 100.0 / len(train_labels))
